export HF_ENDPOINT="https://hf-mirror.com"
export CUDA_VISIBLE_DEVICES=0,1,2,3
NUM_GPUS=4
MODEL=../AnonymousRepo01

export VLLM_WORKER_MULTIPROC_METHOD=spawn
MODEL_ARGS="pretrained=$MODEL,dtype=bfloat16,max_model_length=32768,gpu_memory_utilization=0.8,tensor_parallel_size=$NUM_GPUS"
OUTPUT_DIR=data/evals/$MODEL
TASKS="math_500 aime24 gpqa:diamond"

for TASK in $TASKS
do
   echo "Processing task: $TASK"
lighteval vllm $MODEL_ARGS "custom|$TASK|0|0" \
    --custom-tasks src/open_r1/evaluate.py \
    --use-chat-template \
    --output-dir $OUTPUT_DIR \
    --save-details
done

MODIFIED_STRING="${MODEL//\//_}"
DETAILS_FILEDIRS=$(find $OUTPUT_DIR/details/$MODIFIED_STRING/ -type d \( -name "$(date +%Y-%m-%d*)" \))
if [ -n "$DETAILS_FILEDIRS" ]; then
DETAILS_FILEPATHS=$(find $DETAILS_FILEDIRS -type f \( -name "*.parquet" \))
python parquet2jsonl.py --data_files $DETAILS_FILEPATHS
DETAILS_JSONL_FILEPATHS=$(find $DETAILS_FILEDIRS -type f \( -name "*.jsonl" \))

python test_avg_len.py --files $DETAILS_JSONL_FILEPATHS --tokenizer_path $MODEL
fi